import torch
from torch.utils.data.dataloader import DataLoader
from torch.utils.data.dataset import TensorDataset

from lstm import *
from flask import Flask
from flask import request, session, redirect, jsonify
from utils import getWeekData

app = Flask(__name__, static_folder="static", template_folder='templates')


@app.route("/lstm/pred", methods=["GET"])
def predUser():
	model = LSTM().to(device)
	# 测试模式
	model.load_state_dict(torch.load(lstm_path))
	model.eval()
	userId = request.args.get("userId")
	# 获取近两周的数据用于预测
	weekData = getWeekData(userId)
	train_inout_seq = create_inout_sequences(weekData, 5)
	# 转为二进制字符串
	binary_str = format(int(userId), '020b')
	binary_list = [int(bit) for bit in binary_str]
	binary_tensor = torch.tensor(binary_list).float().view(-1, 20).to(device)

	# 预测数据
	train_data = TensorDataset(torch.FloatTensor([s[0] for s in train_inout_seq]),
	                           torch.FloatTensor([s[1] for s in train_inout_seq]))
	train_loader = DataLoader(train_data, shuffle=False, batch_size=1, drop_last=True)
	pred_list = []
	real_list = []
	for j, (seq, labels) in enumerate(train_loader):
		seq = seq.to(device)
		labels = labels.to(device)
		y_pred = model(seq, binary_tensor)
		real_list.append(labels.item())
		pred_list.append(int(y_pred.item()))
	data = {
		"pred": pred_list,
		"real": real_list
	}
	return jsonify(data)
